#! /bin/python

"""
Fit nonbond parameters by Gromacs
Usage: 
    fit_nonbond.py data_folder
    data_folder
"""


EQUILIBRIUM_MDP = """title           = equilibrium

cpp             = /usr/bin/cpp
integrator      = md
dt              = 0.0010 ; ps !
nsteps          = 50000  ;
nstcomm         = 1
nstxout         = 0
nstvout         = 0
nstfout         = 0
nstlist         = 5
ns_type         = grid
rlist           = 1.0
coulombtype     = Cut-off
rcoulomb        = 1.0
rvdw            = 1.0
DispCorr        = EnerPres
fourierspacing  = 0.12
fourier_nx      = 0
fourier_ny      = 0
fourier_nz      = 0
pme_order       = 4
ewald_rtol      = 1e-5
optimize_fft    = yes

; Temperature control
Tcoupl          = berendsen
tc_grps         = System
tau_t           = 0.1
ref_t           = %(Temperature)f
gen_vel         = yes 
gen_temp        = %(Temperature)f
gen_seed        = 1

; Pressure control
Pcoupl          = berendsen
pcoupltype      = isotropic
tau_p           = 0.2
compressibility = 4.5e-5
ref_p           = 1.01325
"""





SAMPLING_MDP = """title           = sampling

cpp             = /usr/bin/cpp
integrator      = md
dt              = 0.0010 ; ps !
nsteps          = 50000  ;
nstcomm         = 1
nstxout         = 10     ; collect data every 1.0 ps
nstvout         = 0
nstfout         = 0
nstlist         = 5
ns_type         = grid
rlist           = 1.0
coulombtype     = Cut-off
rcoulomb        = 1.0
rvdw            = 1.0
DispCorr        = EnerPres
fourierspacing  = 0.12
fourier_nx      = 0
fourier_ny      = 0
fourier_nz      = 0
pme_order       = 4
ewald_rtol      = 1e-5
optimize_fft    = yes

; Temperature control
Tcoupl          = berendsen
tc_grps         = System
tau_t           = 0.1
ref_t           = %(Temperature)f
gen_vel         = no

; Pressure control 
Pcoupl          = berendsen
pcoupltype      = isotropic
tau_p           = 0.2
compressibility = 4.5e-5
ref_p           = 1.01325
"""

SHELL_COMMANDS = """
grompp -f %(FileName)s.eq.mdp -c %(FileName)s.gro    -p %(FileName)s.top -maxwarn 3 -o %(FileName)s.eq.tpr
mdrun  -s %(FileName)s.eq.tpr -c %(FileName)s.eq.gro -g %(FileName)s.eq.log -o %(FileName)s.eq.trr -e %(FileName)s.eq.edr
grompp -f %(FileName)s.md.mdp -c %(FileName)s.eq.gro -p %(FileName)s.top -maxwarn 3 -o %(FileName)s.md.tpr
mdrun  -s %(FileName)s.md.tpr -c %(FileName)s.md.gro -g %(FileName)s.md.log -o %(FileName)s.md.trr -e %(FileName)s.md.edr
echo Density | g_energy -f %(FileName)s.md.edr
echo Potential | g_energy -f %(FileName)s.md.edr
"""

SINGLE_MOL_SAMPLING_MDP = """title           = sampling

cpp             = /usr/bin/cpp
integrator      = md
dt              = 0.0010 ; ps !
nsteps          = 50000  ;
nstcomm         = 1
nstxout         = 10     ; collect data every 1.0 ps
nstvout         = 0
nstfout         = 0
nstlist         = 5
ns_type         = grid
rlist           = 1.0
coulombtype     = Cut-off
rcoulomb        = 1.0
rvdw            = 1.0
DispCorr        = EnerPres
fourierspacing  = 0.12
fourier_nx      = 0
fourier_ny      = 0
fourier_nz      = 0
pme_order       = 4
ewald_rtol      = 1e-5
optimize_fft    = yes

; Temperature control
Tcoupl          = berendsen
tc_grps         = System
tau_t           = 0.1
ref_t           = %(Temperature)f
gen_vel         = no

"""

SINGLE_MOL_SHELL_COMMANDS = """
grompp -f %(FileName)s.mol.mdp -c %(FileName)s.gro    -p %(FileName)s.top -maxwarn 3 -o %(FileName)s.mol.tpr
mdrun  -s %(FileName)s.mol.tpr -c %(FileName)s.mol.gro -g %(FileName)s.mol.log -o %(FileName)s.mol.trr -e %(FileName)s.mol.edr
echo Potential | g_energy -f %(FileName)s.mol.edr
"""

from os import makedirs
from os.path import isdir, join, abspath
from shutil import rmtree, copy
from kuai.strtool import dump_string
from subprocess import Popen, PIPE, STDOUT

def format_parameters(parameters):
    data = {}
    for i in range(len(parameters)):
        data[str(i)] = parameters[i]
    return forcefield % data
    
def single_mol_potential(model, parameters, FileName):
    work_folder = join(DATA_FOLDER, model.name)
    try:
        if isdir(work_folder):
            rmtree(work_folder)
        makedirs(work_folder)
    except:
        pass
    Temperature = model.temperature
    dump_string(join(work_folder, FileName+'.mol.mdp'), SINGLE_MOL_SAMPLING_MDP % vars())
    copy(join(DATA_FOLDER, FileName + '.top'), join(work_folder, FileName + '.top'))
    copy(join(DATA_FOLDER, FileName + '.gro'), join(work_folder, FileName + '.gro'))
    dump_string(join(work_folder, 'ForceField.itp'), format_parameters(parameters))
    dump_string(join(work_folder, 'job.bat'), SINGLE_MOL_SHELL_COMMANDS % vars())
    pipe = Popen('job.bat', stdin=PIPE, stdout=PIPE, stderr=STDOUT, shell=True, cwd=abspath(work_folder))
    pipe.stdin.close()
    line = pipe.stdout.readline()
    resultP = None
    while line:
        if line.startswith('Potential'):
            resultP = float(line.split()[1])
        line = pipe.stdout.readline()
    pipe.stdout.close()
    pipe.wait()
    del pipe
    return resultP

def calculate(model, parameters, FileName):
    work_folder = join(DATA_FOLDER, model.name)
    try:
        if isdir(work_folder):
            rmtree(work_folder)
        makedirs(work_folder)
    except:
        pass
    Temperature = model.temperature
    dump_string(join(work_folder, FileName+'.eq.mdp'), EQUILIBRIUM_MDP % vars())
    dump_string(join(work_folder, FileName+'.md.mdp'), SAMPLING_MDP % vars())
    copy(join(DATA_FOLDER, FileName + '.top'), join(work_folder, FileName + '.top'))
    copy(join(DATA_FOLDER, FileName + '.gro'), join(work_folder, FileName + '.gro'))
    dump_string(join(work_folder, 'ForceField.itp'), format_parameters(parameters))
    dump_string(join(work_folder, 'job.bat'), SHELL_COMMANDS % vars())
    pipe = Popen('job.bat', stdin=PIPE, stdout=PIPE, stderr=STDOUT, shell=True, cwd=abspath(work_folder))
    pipe.stdin.close()
    line = pipe.stdout.readline()
    resultD, resultP = None, None
    while line:
        if line.startswith('Density'):
            resultD = float(line.split()[2])/1000
        elif line.startswith('Potential'):
            resultP = float(line.split()[1])
        line = pipe.stdout.readline()
    pipe.stdout.close()
    pipe.wait()
    del pipe
    return resultD, resultP
    
class Model:
    pass


buffers = {}

def fitting_function(parameters):
    result = []
    sigma2 = 0
    label = ' '.join(['%8.5f' % i for i in parameters])
    print 'Parameters:', label
    print '%30s%16s%16s%16s' % ('', 'Exp', 'Calc', '%')
    for i in controls:
        labelI = i.name + ':' + label
        if labelI in buffers:
            di, hv = buffers[labelI]
        else:
            h0 = single_mol_potential(i, parameters, i.molname)
            di, hv = calculate(i, parameters, i.name)
            if h0 is None or hv is None:
                pass
            else:
                hv = h0 - hv / i.nmol + 8.314e-3 * i.temperature
            buffers[labelI] = di, hv
        
        if di is None:
            dd = 100
            print '%30s%16.5f%16s%16.5f' % (i.name + '  [D]', i.density, 'N/A', dd * 100)
        else:
            dd = (di-i.density) / i.density
            print '%30s%16.5f%16.5f%16.5f' % (i.name + '  [D]', i.density, di, dd * 100)
        if hv is None:
            dh = 100
            print '%30s%16.5f%16s%16.5f' % (i.name + ' [Hv]', i.hv, 'N/A', dh * 100)
        else:
            dh = (hv-i.hv) / i.hv
            print '%30s%16.5f%16.5f%16.5f' % (i.name + ' [Hv]', i.hv, hv, dh * 100)
        result += [dd, dh]
        sigma2 += dd*dd + dh*dh
    print "Sigma2 = ", sigma2
    print
    return result

def sigma2(parameters):
    result = 0
    v = fitting_function(parameters)
    for i in v:
        result += i*i
    return result

if __name__ == '__main__':
    from sys import argv
    from kuai.strtool import create_array, load_string, load_value
    from datetime import datetime
    if len(argv) != 2 or not isdir(argv[1]):
        print __doc__
    else:
        time1 = datetime.today()
        print 'Job started at', time1
        DATA_FOLDER = argv[1]
        datafile = join(DATA_FOLDER, 'Data.txt')
        forcefield = load_string(join(DATA_FOLDER, 'Nonbond.itp'))
        with open(datafile) as file:
            models = create_array(file, Model)

        execfile(join(DATA_FOLDER, 'NonbondControl.py'))
        
        models_to_fit = []
        for i in models:
            if i.name in controls:
                models_to_fit.append(i)
        controls = models_to_fit

        from scipy.optimize import leastsq
        result = leastsq(fitting_function, parameters, full_output=True, col_deriv=True, ftol=0.005, xtol=0.005, epsfcn=0.005)

        print '\n\n\n\n', 'The final result is:', result
        dump_string(join(DATA_FOLDER, 'Result.itp'), format_parameters(result[0]))

        time2 = datetime.today()
        print 'Job finished at', time2
        print 'It costed', time2-time1
   